import asyncio
import datetime
import json
import os


import uuid

from sqlalchemy.orm import Session
from fastapi import Depends, WebSocket
import models
from models import Share
from schemas import response_success, response_error
from models.db import get_db
from fastapi import APIRouter

from schemas.toolkit import ToolkitDebugSchema
from views import generate_script, oauth2_scheme, get_current_user_name

router = share_router = APIRouter()


@router.post('')
def share(form: ToolkitDebugSchema, db: Session = Depends(get_db), token: str = Depends(oauth2_scheme)):
    script = db.query(models.Toolkit).filter_by(id=form.id).first()
    if not script:
        return response_error('数据不存在')

    params = {x['id']: x for x in json.loads(script.args) if x.get('private') is False}
    for k, v in form.args.items():
        if k in params:
            params[k]['value'] = v

    username = get_current_user_name(token)
    share_id = uuid.uuid4().hex
    db.add(Share(name=script.name, share_id=share_id, script_id=script.id, args=json.dumps(list(params.values())),
                 content=form.content, username=username))
    db.commit()
    return response_success('ok', data={'share_id': share_id})


@router.get('')
def share_detail(share_id: str, db: Session = Depends(get_db)):
    db_share = db.query(Share).filter_by(share_id=share_id).first()
    if not db_share:
        return response_error('分享不存在')
    db_share.args = json.loads(db_share.args)
    return response_success('ok', data={'detail': db_share})


@router.websocket('/start')
async def start(websocket: WebSocket, db: Session = Depends(get_db)):
    import pty
    import pwd
    await websocket.accept()
    params = await websocket.receive_json()
    db_share = db.query(models.Share).filter_by(share_id=params.get('share_id')).first()
    if not db_share:
        await websocket.send_text('分享不存在 ! ...\r\n')
        return
    db_script = db.query(models.Toolkit).filter_by(id=db_share.script_id).first()
    if not db_script:
        await websocket.send_text('脚本不存在 ! ...\r\n')
        return

    if not db_share.content and not db_script.content:
        await websocket.send_text('没有代码可供执行 ! ...\r\n')
        return

    command = generate_script(db_script, {'args': json.loads(db_share.args)})
    workspace = db_script.workspace
    exec_user = db_script.exec_user

    pid, fd = pty.fork()
    if pid == 0:
        if workspace:
            os.chdir(workspace)
        if exec_user:
            os.setuid(pwd.getpwnam(exec_user).pw_uid)
        # 在子进程中执行命令.
        os.execvp(command[0], command)
    else:
        await websocket.send_json({
            'end': False,
            'pid': pid,
        })
        history = b''
        # 在父进程中读取输出并显示
        try:
            while True:
                try:
                    read_bytes = os.read(fd, 1024)
                    if not read_bytes:
                        break
                    history += read_bytes
                    await websocket.send_bytes(read_bytes)
                    await asyncio.sleep(0)
                except OSError:
                    break
        finally:
            # 确保文件描述符被关闭
            os.close(fd)
        # 等待子进程结束，并获取其退出状态
    _, exit_status = os.waitpid(pid, 0)

    await websocket.send_json({
        'end': True,
        'status': exit_status,
    })

    await websocket.close()
    toolkit = models.ToolkitLog(tool_id=db_script.id,
                                tool_name=db_script.name,
                                create_time=datetime.datetime.now(),
                                username=db_share.username,
                                status=exit_status,
                                content=history.decode())
    db.add(toolkit)
    db.commit()

    return response_success('ok', data=db_share)
